// Copyright 2022 jeff.li. and/or its affiliates.

#include <tbir/pipeline/device_op.h>
#include <tbir/runtime/container/ndarray_helper.h>
#include "tbir/pipeline/attributes.h"
#include "tbir/pipeline/internal_helper_funcs.h"
#include "tbir/pipeline/tx_session.h"
#include "tbir/runtime/c_runtime_api.h"
#include "tbir/runtime/device_api.h"
#include "tbir/runtime/dlpack.h"
#include "tbir/runtime/logging.h"
#include "tbir/runtime/threadpool/lock_based_thread_pool.h"

namespace tbir::runtime {

    void DeviceOp::Init() {
        auto device = GetAttr<Unicode>("device");
        auto dl_dev = NDArrayHelper::GetDevice(device);
        if (dl_dev.device_type == kDLCPU) {
            device_id_ = -1;
        } else {
            device_id_ = dl_dev.device_id;
        }
    }

    RTValue DeviceOp::Process(PyArgs inputs) const {
        int session_device_id = device_;
        TbirStreamHandle current_stream = nullptr;
        void *thread_pool = nullptr;

        if (device_ == NONE_DEVICE) {
            // no session, use global
            if (device_id_ >= 0) {
                TbirDevice ctx{kDLCUDA, device_id_};
                DeviceAPI *api = DeviceAPI::Get(ctx, true);
                if (api != nullptr) {
                    current_stream = api->GetCurrentThreadStream(ctx);
                }
            }
        } else {
            if (device_ >= 0) {
                session_device_id = internal::cuda_device_offset(device_);
                TbirDevice ctx{kDLCUDA, session_device_id};
                DeviceAPI *api = DeviceAPI::Get(ctx, true);
                if (api != nullptr) {
                    current_stream = api->GetCurrentThreadStream(ctx);
                }
            }
        }
        thread_pool = belong_to_->GetComputeThreadPool();
        // TODO: remove h2d and d2h
        return Dict({{"device_id",         device_id_},
                     {"session_device_id", session_device_id},
                     {"compute_stream",    current_stream},
                     {"h2d_stream",        current_stream},
                     {"d2h_stream",        current_stream},
                     {"thread_pool",       thread_pool}});
    }

    // Device should be rebind for every session
    TBIR_REGISTER_NATIVE_OP(DeviceOp).SetThreadSafety(false);

}  // namespace tbir::runtime